{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nb90LCrotIx4"
      },
      "source": [
        "## Cellpose-SAM: superhuman generalization for cellular segmentation\n",
        "\n",
        "Marius Pachitariu, Michael Rariden, Carsen Stringer\n",
        "\n",
        "[paper](https://www.biorxiv.org/content/10.1101/2025.04.28.651001v1) | [code](https://github.com/MouseLand/cellpose)\n",
        "\n",
        "This notebook shows how to process your own 2D or 3D images, saved on Google Drive.\n",
        "\n",
        "This notebook is adapted from the notebook by Pradeep Rajasekhar, inspired by the [ZeroCostDL4Mic notebook series](https://github.com/HenriquesLab/ZeroCostDL4Mic/wiki)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z0s2fz5hUk75"
      },
      "source": [
        "### Make sure you have GPU access enabled by going to Runtime -> Change Runtime Type -> Hardware accelerator and selecting GPU\n",
        "\n",
        "![image.png]()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VZe3EmLeDr5h"
      },
      "source": [
        "### Mount your google drive\n",
        "\n",
        "Run this cell to connect your Google Drive to colab:\n",
        "* Click on the URL.\n",
        "* Sign in your Google Account.\n",
        "\n",
        "You will either have to:\n",
        "* copy the authorisation code and enter it into box below OR\n",
        "* in the new google colab, you can just click \"Allow\" and it should connect."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b8wPXz7PGNLt"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l9YpMcAuDr5i"
      },
      "source": [
        "\n",
        "Then click on \"Folder\" icon on the Left, press the refresh button. Your Google Drive folder should now be available here as \"gdrive\".\n",
        "\n",
        "\n",
        "![image.png]()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yH4zTaWwvH9x"
      },
      "source": [
        "Click on the triangle icon and it will allow you to access whole drive. Navigate to the folder containing your images. Once you are there, click on the three dots on the right of the folder and select \"Copy Path\"\n",
        "\n",
        "![image.png]()\n",
        "\n",
        "Copy and paste this path in the **dir** string below"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h_iAN7cAthma"
      },
      "source": [
        "### Install Cellpose-SAM\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hG3LSmJmLylT"
      },
      "outputs": [],
      "source": [
        "!pip install git+https://www.github.com/mouseland/cellpose.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fGSKx-EiDr5k"
      },
      "source": [
        "Check GPU and instantiate model - will download weights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "5ydQ-fggSiUm"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from cellpose import models, core, io, plot\n",
        "from pathlib import Path\n",
        "from tqdm import trange\n",
        "import matplotlib.pyplot as plt\n",
        "from natsort import natsorted\n",
        "\n",
        "io.logger_setup() # run this to get printing of progress\n",
        "\n",
        "#Check if colab notebook instance has GPU access\n",
        "if core.use_gpu()==False:\n",
        "  raise ImportError(\"No GPU access, change your runtime\")\n",
        "\n",
        "model = models.CellposeModel(gpu=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vOpZWZtDDr5l"
      },
      "source": [
        "Input directory with your images:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-lZP6alpUAfY"
      },
      "outputs": [],
      "source": [
        "# *** change to your google drive folder path ***\n",
        "dir = \"/content/gdrive/MyDrive/PATH-TO-FILES/\"\n",
        "dir = Path(dir)\n",
        "if not dir.exists():\n",
        "  raise FileNotFoundError(\"directory does not exist\")\n",
        "\n",
        "# *** change to your image extension ***\n",
        "image_ext = \".tif\"\n",
        "\n",
        "# list all files\n",
        "files = natsorted([f for f in dir.glob(\"*\"+image_ext) if \"_masks\" not in f.name and \"_flows\" not in f.name])\n",
        "\n",
        "if(len(files)==0):\n",
        "  raise FileNotFoundError(\"no image files found, did you specify the correct folder and extension?\")\n",
        "else:\n",
        "  print(f\"{len(files)} images in folder:\")\n",
        "\n",
        "for f in files:\n",
        "  print(f.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eyj9hnQvDr5l"
      },
      "source": [
        "## Run Cellpose-SAM on one image in folder\n",
        "\n",
        "Here are some of the parameters you can change:\n",
        "\n",
        "* ***flow_threshold*** is  the  maximum  allowed  error  of  the  flows  for  each  mask.   The  default  is 0.4.\n",
        "    *  **Increase** this threshold if cellpose is not returning as many masks as you’d expect (or turn off completely with 0.0)\n",
        "    *   **Decrease** this threshold if cellpose is returning too many ill-shaped masks.\n",
        "\n",
        "* ***cellprob_threshold*** determines proability that a detected object is a cell.   The  default  is 0.0.\n",
        "    *   **Decrease** this threshold if cellpose is not returning as many masks as you’d expect or if masks are too small\n",
        "    *   **Increase** this threshold if cellpose is returning too many masks esp from dull/dim areas.\n",
        "\n",
        "* ***tile_norm_blocksize*** determines the size of blocks used for normalizing the image. The default is 0, which means the entire image is normalized together.\n",
        "  You may want to change this to 100-200 pixels if you have very inhomogeneous brightness across your image.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vefSZf0tLxKI"
      },
      "outputs": [],
      "source": [
        "img = io.imread(files[0])\n",
        "\n",
        "print(f'your image has shape: {img.shape}. Assuming channel dimension is last with {img.shape[-1]} channels')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CfhdXrkmMV2p"
      },
      "source": [
        "### Channel Selection: \n",
        "\n",
        "- Use the dropdowns below to select the _zero-indexed_ channels of your image to segment. The order does not matter. Remember to rerun the cell after you edit the dropdowns.\n",
        "\n",
        "- If you have a histological image taken in brightfield, you don't need to adjust the channels.\n",
        "\n",
        "- If you have a fluroescent image with multiple stains, you should choose one channel with a cytoplasm/membrane stain, one channel with a nuclear stain, and set the third channel to `None`. Choosing multiple channels may produce segmentaiton of all the structures in the image. If you have retrained the model on your data with a thrid stain (described below), you can run segmentation with all channels. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8a4oOuFZKoA-"
      },
      "outputs": [],
      "source": [
        "first_channel = '0' # @param ['None', 0, 1, 2, 3, 4, 5]\n",
        "second_channel = '1' # @param ['None', 0, 1, 2, 3, 4, 5]\n",
        "third_channel = '2' # @param ['None', 0, 1, 2, 3, 4, 5]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "98DA8zm4A__m"
      },
      "outputs": [],
      "source": [
        "selected_channels = []\n",
        "for i, c in enumerate([first_channel, second_channel, third_channel]):\n",
        "  if c == 'None':\n",
        "    continue\n",
        "  if int(c) > img.shape[-1]:\n",
        "    assert False, 'invalid channel index, must have index greater or equal to the number of channels'\n",
        "  if c != 'None':\n",
        "    selected_channels.append(int(c))\n",
        "\n",
        "\n",
        "\n",
        "img_selected_channels = np.zeros_like(img)\n",
        "img_selected_channels[:, :, :len(selected_channels)] = img[:, :, selected_channels]\n",
        "\n",
        "\n",
        "flow_threshold = 0.4\n",
        "cellprob_threshold = 0.0\n",
        "tile_norm_blocksize = 0\n",
        "\n",
        "masks, flows, styles = model.eval(img_selected_channels, batch_size=32, flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold,\n",
        "                                  normalize={\"tile_norm_blocksize\": tile_norm_blocksize})\n",
        "\n",
        "fig = plt.figure(figsize=(12,5))\n",
        "plot.show_segmentation(fig, img_selected_channels, masks, flows[0])\n",
        "plt.tight_layout()\n",
        "plt.show()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3JRxBPmatrK7"
      },
      "source": [
        "## Run Cellpose-SAM on folder of images\n",
        "\n",
        "if you have many large images, you may want to run them as a loop over images\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JgfGlUJNDr5n"
      },
      "outputs": [],
      "source": [
        "masks_ext = \".png\" if image_ext == \".png\" else \".tif\"\n",
        "for i in trange(len(files)):\n",
        "    f = files[i]\n",
        "    img = io.imread(f)\n",
        "    masks, flows, styles = model.eval(img, batch_size=32, flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold,\n",
        "                                  normalize={\"tile_norm_blocksize\": tile_norm_blocksize})\n",
        "    io.imsave(dir / (f.stem + \"_masks\" + masks_ext), masks)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZsmgS0KJDr5n"
      },
      "source": [
        "if you have small images, you may want to load all of them first and then run, so that they can be batched together on the GPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XH-K72hGDr5n"
      },
      "outputs": [],
      "source": [
        "print(\"loading images\")\n",
        "imgs = [io.imread(files[i]) for i in trange(len(files))]\n",
        "\n",
        "print(\"running cellpose-SAM\")\n",
        "masks, flows, styles = model.eval(imgs, batch_size=32, flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold,\n",
        "                                  normalize={\"tile_norm_blocksize\": tile_norm_blocksize})\n",
        "\n",
        "print(\"saving masks\")\n",
        "for i in trange(len(files)):\n",
        "    f = files[i]\n",
        "    io.imsave(dir / (f.stem + \"_masks\" + masks_ext), masks[i])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9GdeH0nCDr5n"
      },
      "source": [
        "to save your masks for ImageJ, run the following code:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "s7NAuncdY618"
      },
      "outputs": [],
      "source": [
        "for i in trange(len(files)):\n",
        "    f = files[i]\n",
        "    masks0 = io.imsave(dir / (f.name + \"_masks\" + masks_ext))\n",
        "    io.save_rois(masks0, f)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "cellpose",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
